from collections import defaultdict

from flask import Flask, request, Response
from TTS.api import TTS
from utils import create_wav_header
import io
import soundfile as sf
from TTS.utils.radam import RAdam
import torch.serialization
app = Flask(__name__)
# 加入 RAdam 到允许的安全加载列表
torch.serialization.add_safe_globals([RAdam,defaultdict,dict])

# 初始化 TTS（可改为中文模型）
tts = TTS(model_name="tts_models/zh-CN/baker/tacotron2-DDC-GST", progress_bar=False, gpu=False)

@app.route('/tts', methods=['GET'])
def tts_stream():
    text = request.args.get('text', '你好，请提供一个文本参数 text')
    sample_rate = 22050

    def generate():
        # WAV header
        yield create_wav_header(sample_rate=sample_rate)

        # 生成音频数据
        audio = tts.tts(text)
        buffer = io.BytesIO()
        sf.write(buffer, audio, sample_rate, format='WAV')
        buffer.seek(44)  # 跳过 WAV header 部分（我们已经手动生成了）
        while True:
            chunk = buffer.read(1024)
            if not chunk:
                break
            yield chunk

    return Response(generate(), mimetype="audio/wav")

if __name__ == '__main__':
    app.run(debug=True)